from flask import Flask, request, jsonify
from flask_cors import CORS  # 引入 CORS
import pymysql
import random
import smtplib
from email.mime.text import MIMEText
from email.header import Header
import hashlib
import time
import jwt  # 引入 PyJWT
import datetime  # 引入 datetime 用于设置 token 过期时间

app = Flask(__name__)

# 配置 CORS，允许所有来源访问（生产环境可限制特定域名）
CORS(app, resources={r"/*": {"origins": "*"}})

# JWT 密钥（生产环境建议使用环境变量存储）
SECRET_KEY = "your-secret-key"  # 请替换为一个安全的密钥

# 数据库连接配置
def get_db_connection():
    return pymysql.connect(
        host='localhost',
        user='root',  # 替换为您的 MySQL 用户名
        password='123456',  # 替换为您的 MySQL 密码
        database='dachuang',
        charset='utf8mb4',
        cursorclass=pymysql.cursors.DictCursor
    )

# 生成6位验证码
def generate_verification_code():
    return ''.join([str(random.randint(0, 9)) for _ in range(6)])

# 发送QQ邮箱验证码
def send_verification_email(email, code):
    sender = '952304392@qq.com'  # 替换为您的QQ邮箱
    password = 'lwlbpkrftgzxbbfb'  # 替换为您的QQ邮箱授权码

    msg = MIMEText(f'您的验证码是：{code}，有效期20分钟', 'plain', 'utf-8')  # 更新提示信息
    msg['From'] = Header(sender)
    msg['To'] = Header(email)
    msg['Subject'] = Header('验证码')

    try:
        smtp_obj = smtplib.SMTP_SSL('smtp.qq.com', 465)
        smtp_obj.login(sender, password)
        smtp_obj.sendmail(sender, [email], msg.as_string())
        smtp_obj.quit()
        return True
    except Exception as e:
        print(f"邮件发送失败: {e}")
        return False

# 密码哈希
def hash_password(password):
    return hashlib.sha256(password.encode()).hexdigest()

# 生成 JWT Token
def generate_token(username):
    payload = {
        'username': username,
        'exp': datetime.datetime.utcnow() + datetime.timedelta(hours=24)  # Token 有效期 24 小时
    }
    return jwt.encode(payload, SECRET_KEY, algorithm='HS256')

# 验证 JWT Token
def verify_token(token):
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=['HS256'])
        return payload['username']
    except jwt.ExpiredSignatureError:
        return None  # Token 已过期
    except jwt.InvalidTokenError:
        return None  # Token 无效

# 请求注册验证码
@app.route('/register/code', methods=['POST'])
def register_code():
    data = request.get_json()
    username = data.get('username')
    email = data.get('email')

    if not all([username, email]):
        return jsonify({'message': '缺少必要参数'}), 400

    conn = get_db_connection()
    c = conn.cursor()

    c.execute('SELECT * FROM users WHERE username = %s OR email = %s', (username, email))
    if c.fetchone():
        conn.close()
        return jsonify({'message': '用户名或邮箱已存在'}), 400

    code = generate_verification_code()

    if send_verification_email(email, code):
        c.execute('INSERT INTO verification_codes (email, code, timestamp) VALUES (%s, %s, %s) '
                  'ON DUPLICATE KEY UPDATE code = %s, timestamp = %s',
                  (email, code, int(time.time()), code, int(time.time())))
        conn.commit()
        conn.close()
        return jsonify({'message': '验证码已发送至您的邮箱'})
    else:
        conn.close()
        return jsonify({'message': '验证码发送失败'}), 500

# 注册接口
@app.route('/register', methods=['POST'])
def register():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')
    email = data.get('email')
    code = data.get('code')

    if not all([username, password, email, code]):
        return jsonify({'message': '缺少必要参数'}), 400

    conn = get_db_connection()
    c = conn.cursor()

    c.execute('SELECT code, timestamp FROM verification_codes WHERE email = %s', (email,))
    result = c.fetchone()

    if not result:
        conn.close()
        return jsonify({'message': '验证码无效'}), 400

    stored_code, timestamp = result['code'], result['timestamp']

    if int(time.time()) - timestamp > 1200:  # 改为 20 分钟（1200秒）
        conn.close()
        return jsonify({'message': '验证码已过期'}), 400

    if stored_code != code:
        conn.close()
        return jsonify({'message': '验证码错误'}), 400

    hashed_password = hash_password(password)
    c.execute('INSERT INTO users (username, password, email, verified) VALUES (%s, %s, %s, 1)',
              (username, hashed_password, email))
    c.execute('DELETE FROM verification_codes WHERE email = %s', (email,))
    conn.commit()
    conn.close()

    return jsonify({'message': '注册成功'})

# 请求登录验证码
@app.route('/login/code', methods=['POST'])
def login_code():
    data = request.get_json()
    email = data.get('email')

    if not email:
        return jsonify({'message': '缺少邮箱参数'}), 400

    conn = get_db_connection()
    c = conn.cursor()

    c.execute('SELECT * FROM users WHERE email = %s', (email,))
    if not c.fetchone():
        conn.close()
        return jsonify({'message': '邮箱未注册'}), 400

    code = generate_verification_code()

    if send_verification_email(email, code):
        c.execute('INSERT INTO verification_codes (email, code, timestamp) VALUES (%s, %s, %s) '
                  'ON DUPLICATE KEY UPDATE code = %s, timestamp = %s',
                  (email, code, int(time.time()), code, int(time.time())))
        conn.commit()
        conn.close()
        return jsonify({'message': '验证码已发送至您的邮箱'})
    else:
        conn.close()
        return jsonify({'message': '验证码发送失败'}), 500

# 登录接口（添加 token）
@app.route('/login', methods=['POST'])
def login():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')
    email = data.get('email')
    code = data.get('code')

    if not all([username, password, email, code]):
        return jsonify({'message': '缺少必要参数'}), 400

    conn = get_db_connection()
    c = conn.cursor()

    c.execute('SELECT password, email FROM users WHERE username = %s', (username,))
    result = c.fetchone()

    if not result:
        conn.close()
        return jsonify({'message': '用户不存在'}), 400

    stored_password, stored_email = result['password'], result['email']

    if stored_email != email:
        conn.close()
        return jsonify({'message': '邮箱不匹配'}), 400

    if hash_password(password) != stored_password:
        conn.close()
        return jsonify({'message': '密码错误'}), 400

    c.execute('SELECT code, timestamp FROM verification_codes WHERE email = %s', (email,))
    result = c.fetchone()

    if not result:
        conn.close()
        return jsonify({'message': '验证码无效'}), 400

    stored_code, timestamp = result['code'], result['timestamp']

    if int(time.time()) - timestamp > 1200:  # 改为 20 分钟（1200秒）
        conn.close()
        return jsonify({'message': '验证码已过期'}), 400

    if stored_code != code:
        conn.close()
        return jsonify({'message': '验证码错误'}), 400

    c.execute('DELETE FROM verification_codes WHERE email = %s', (email,))
    conn.commit()
    conn.close()

    # 生成 token 并返回
    token = generate_token(username)
    return jsonify({'message': '登录成功', 'token': token})

# 示例受保护接口（需要 token）
@app.route('/protected', methods=['GET'])
def protected():
    token = request.headers.get('Authorization')
    if not token:
        return jsonify({'message': '缺少 token'}), 401
    if token.startswith('Bearer '):
        token = token[7:]  # 移除 "Bearer " 前缀
    username = verify_token(token)
    if not username:
        return jsonify({'message': '无效或过期的 token'}), 401
    return jsonify({'message': f'欢迎, {username}!'})

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)